const express = require('express');
const router = express.Router();
const { spawn } = require('child_process');
const path = require('path');
const fs = require('fs');
const WebSocket = require('ws');

// 创建WebSocket服务器
const wss = new WebSocket.Server({ port: 3457 });
let activeTrainingProcess = null;

// 广播训练进度给所有连接的客户端
const broadcast = (data) => {
    wss.clients.forEach(client => {
        if (client.readyState === WebSocket.OPEN) {
            client.send(JSON.stringify(data));
        }
    });
};

// 获取已上传的文件列表
router.get('/files', async (req, res) => {
    try {
        const uploadDir = path.join(__dirname, '../upload');
        // 确保目录存在
        if (!fs.existsSync(uploadDir)) {
            fs.mkdirSync(uploadDir, { recursive: true });
        }

        // 读取目录中的所有文件
        const files = fs.readdirSync(uploadDir)
            .filter(file => file.toLowerCase().endsWith('.csv'))
            .map((filename) => ({
                filename: filename,
                path: path.join(uploadDir, filename)
            }));

        res.json({
            success: true,
            files: files
        });
    } catch (error) {
        console.error('Error getting file list:', error);
        res.status(500).json({
            success: false,
            message: error.message
        });
    }
});

// 开始训练
router.post('/train', async (req, res) => {
    try {
        const { topology, trainFilePath, routeFilePath, params } = req.body;

        // 验证必要参数
        if (!topology || !trainFilePath || !routeFilePath || !params) {
            return res.status(400).json({
                success: false,
                message: '缺少必要参数'
            });
        }

        // 如果已有训练进程在运行，返回错误
        if (activeTrainingProcess) {
            return res.status(400).json({
                success: false,
                message: '已有训练任务在进行中'
            });
        }

        // 验证文件是否存在
        if (!fs.existsSync(trainFilePath)) {
            return res.status(400).json({
                success: false,
                message: '训练文件不存在'
            });
        }
        if (!fs.existsSync(routeFilePath)) {
            return res.status(400).json({
                success: false,
                message: '路由矩阵文件不存在'
            });
        }

        // 设置训练参数
        const trainScript = path.join(__dirname, '../../../../FlowTM-main/train.py');
        const outputDir = path.join(__dirname, '../../../../FlowTM-main/OUTPUT');
        const dataDir = path.join(__dirname, '../../../../FlowTM-main/Data');

        // 确保输出目录存在
        if (!fs.existsSync(outputDir)) {
            fs.mkdirSync(outputDir, { recursive: true });
        }

        // 确保数据目录存在
        if (!fs.existsSync(dataDir)) {
            fs.mkdirSync(dataDir, { recursive: true });
        }

        // 复制文件到数据目录
        const trainFileName = path.basename(trainFilePath);
        const routeFileName = path.basename(routeFilePath);
        const trainTargetPath = path.join(dataDir, trainFileName);
        const routeTargetPath = path.join(dataDir, routeFileName);

        fs.copyFileSync(trainFilePath, trainTargetPath);
        fs.copyFileSync(routeFilePath, routeTargetPath);

        const args = [
            trainScript,
            '--data-root', dataDir,
            '--dataset', topology === 'topology1' ? 'abilene' : 'geant',
            '--rm-fn', routeTargetPath,
            '--train-epoch', params.epochs.toString(),
            '--lr', params.learningRate.toString(),
            '--batch-size', params.batchSize.toString(),
            '--run-name', outputDir
        ];

        console.log('Starting training with args:', args.join(' '));

        // 启动训练脚本
        activeTrainingProcess = spawn('python', args);

        // 处理输出
        activeTrainingProcess.stdout.on('data', (data) => {
            const message = data.toString();
            console.log('Training output:', message);
            // 直接转发训练脚本的输出
            broadcast({ type: 'log', message });
        });

        activeTrainingProcess.stderr.on('data', (data) => {
            const message = data.toString();
            console.error('Training error:', message);
            // 直接转发错误输出
            broadcast({ type: 'log', message });
        });

        activeTrainingProcess.on('close', (code) => {
            console.log('Training process ended with code:', code);
            broadcast({ type: 'log', message: `训练进程已结束，退出码: ${code}` });
            activeTrainingProcess = null;
        });

        res.json({ success: true, message: '训练已启动' });
    } catch (error) {
        console.error('Training error:', error);
        res.status(500).json({ success: false, message: error.message });
    }
});

// 停止训练
router.post('/train/stop', async (req, res) => {
    try {
        if (activeTrainingProcess) {
            activeTrainingProcess.kill();
            activeTrainingProcess = null;
            res.json({ success: true, message: '训练已停止' });
        } else {
            res.status(400).json({ success: false, message: '没有正在进行的训练任务' });
        }
    } catch (error) {
        res.status(500).json({ success: false, message: error.message });
    }
});

module.exports = router; 